import asyncio
import json
import websockets
import socket
import requests
from qiniu import Auth, put_file_v2, etag
from datetime import datetime
import mss
import configparser
import cv2
import os


class QNStore():
    def __init__(self, access_key, secret_key, bucket_name, base_url):
        self.auth = Auth(access_key, secret_key)
        self.bucket_name = bucket_name
        self.base_url = base_url

    def get_url(self, filename):
        return self.auth.private_download_url(f"http://{self.base_url}/{filename}", expires=3600)

    def upload(self, filename, localfile):
        token = self.auth.upload_token(self.bucket_name, filename, 3600)
        ret, _ = put_file_v2(token, filename, localfile, version='v2')
        
        if ret['hash'] == etag(localfile):
            return self.get_url(filename)
        else:
            return None
        
        
class AiRequest():
    def __init__(self, url, model, api_key):
        self.model = model
        self.url = url
        self.headers = {
            "Authorization": f"Bearer {api_key}", 
            "Content-Type": "application/json"
        }
        self.payload = {
            "model": model,
            "stream": False,
            "max_tokens": 2048,
            "messages": []
        }

    def chat(self, img_link):
        message = [
            {
                "role": "system", 
                "content": "你是一个在图像中寻找试卷的题目并给出答案的助手，是一个不需要复述问题和过程的简洁助手。"
            },
            {   "role": "user", 
                    "content": [
                        {"type": "text", "text": "请查看下面的图片并简短地回答其中的题目，不需要解析与过程。"},
                        {"type": "image_url", "image_url": {"url": img_link}}
                    ]
            }
        ]
        payload = self.payload.copy()
        payload['messages'] = message
        response = requests.post(self.url, headers=self.headers, json=payload)
        return response.text
    
    def code(self, img_link):
        message = [
            {
                "role": "system", 
                "content": "你是一个在图像中寻找试卷的题目并给出答案的助手，是一个不需要复述问题和过程的简洁助手。"
            },
            {   "role": "user", 
                    "content": [
                        {"type": "text", "text": "图片中包含一道算法题，找出这道算法题并给出完成该题的C++的代码。"},
                        {"type": "image_url", "image_url": {"url": img_link}}
                    ]
            }
        ]
        payload = self.payload.copy()
        payload['messages'] = message
        response = requests.post(self.url, headers=self.headers, json=payload)
        return response.text
    

class ScreenCapture():
    def __init__(self):
        with mss.mss() as sct:
            if len(sct.monitors) - 1 > 1:
                print("\033[91mMultiple monitors detected. Using the primary monitor.\033[0m")

    def capture(self) -> str:
        try:
            with mss.mss() as sct:
                timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
                filename = f"SCR_{timestamp}.png"
                save_dir = os.path.join(".", "img_cache", filename)
                sct.shot(output=save_dir)
                return filename
        except Exception as e:
            return f"Error: {str(e)}"
        

class CamCapture():
    def __init__(self, index=0, width=640, height=480, lrFlip=False, udFlip=False):
        self.index = index
        self.width = width
        self.height = height
        self.lrFlip = lrFlip
        self.udFlip = udFlip
        self.lock = False
    
    def capture(self) -> str:
        if self.lock:
            return "Error: Camera is busy."
        self.lock = True
        cap = cv2.VideoCapture(self.index)
        if not cap.isOpened():
            return f"Error: Could not open {self.index} camera."
        try:
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
            real_width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
            real_height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
            if real_width != self.width or real_height != self.height:
                print(f"Warning: Could not set camera resolution to {self.width}x{self.height}, actual is {real_width}x{real_height}")
            for _ in range(5):
                ret, frame = cap.read()
            ret, frame = cap.read()
            if not ret:
                return f"Error: Could not read frame from camera."
            timestamp = datetime.now().strftime("%Y-%m-%d_%H%M%S")
            filename = f"CAM_{timestamp}.png"
            if self.lrFlip:
                frame = cv2.flip(frame, 1)
            if self.udFlip:
                frame = cv2.flip(frame, 0)
            save_dir = os.path.join(".", "img_cache", filename)
            cv2.imwrite(save_dir, frame)
            cap.release()
            return filename
        except Exception as e:
            cap.release()
            return f"Error: {str(e)}"
        finally:
            cap.release()
            self.lock = False


def process_command(command):
    if (len(command) > 0):
        filename = capture.capture()
        if filename.startswith("Err"):
            filename = capture.capture()
            if filename.startswith("Err"):
                return json.dumps({"error": filename})
        save_dir = os.path.join(".", "img_cache", filename)
        img_link = qn_store.upload(filename, save_dir)
        if img_link is None:
            return json.dumps({"error": "Failed to upload image."})
        if command[0] == "2":
            response = ai_request.code(img_link)
        else:
            response = ai_request.chat(img_link)
        return response
    else:
        return json.dumps({"error": "Invalid command format."})


async def echo_handler(websocket):
    client_address = websocket.remote_address
    print(f"client {client_address} connected")
    try:
        async for message in websocket:
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            print(f"[{timestamp}] Recv from {client_address}: {message}")
            try:
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(None, process_command, message)
                await websocket.send(response)
            except Exception as e:
                error_msg = json.dumps({"error": str(e)})
                await websocket.send(error_msg)
                print(f"[{timestamp}] Error processing request from {client_address}: {error_msg}")
            
            
    except websockets.exceptions.ConnectionClosed:
        print(f"client {client_address} disconnected")


async def main():
    host = socket.gethostname()
    print(f"Your Server IP addresses: {host}")
    address = cfg.get('server', 'host')
    print(f"Starting WebSocket server on ws://{address}:8765")
    server = await websockets.serve(echo_handler, address, 8765)
    await server.wait_closed()


if __name__ == "__main__":
    cfg = configparser.ConfigParser()
    if cfg.read('./config.ini', encoding='utf-8') == []:
        print("Error: Could not read config.ini. Please create it based on cfg_details.ini.")
        exit(1)

    ACCESS_KEY = cfg.get('qiniu', 'access')
    SECRET_KEY = cfg.get('qiniu', 'secret')
    BUCKET_NAME = cfg.get('qiniu', 'bucket')
    BASE_URL = cfg.get('qiniu', 'base_url')

    AI_URL = cfg.get('ai', 'url')
    AI_MODEL = cfg.get('ai', 'model')
    API_KEY = cfg.get('ai', 'api_key')

    CAP_MODE = cfg.getint('cap', 'mode')
    if CAP_MODE == 1:
        capture = ScreenCapture()
    else:   
        CAM_INDEX = cfg.getint('cam', 'index')
        CAM_WIDTH = cfg.getint('cam', 'width')
        CAM_HEIGHT = cfg.getint('cam', 'height')
        CAM_LRFLIP = cfg.getboolean('cam', 'lrFlip')
        CAM_UDFLIP = cfg.getboolean('cam', 'udFlip')
        capture = CamCapture(CAM_INDEX, CAM_WIDTH, CAM_HEIGHT, CAM_LRFLIP, CAM_UDFLIP)

    qn_store = QNStore(ACCESS_KEY, SECRET_KEY, BUCKET_NAME, BASE_URL)
    ai_request = AiRequest(AI_URL, AI_MODEL, API_KEY)
    

    asyncio.run(main())
